import codecs

MOD = 256


'''
Key Scheduling Algorithm
'''
def KSA(key):
    
    key_len = len(key)
    S = list(range(MOD))
    
    j = 0
    for i in range(MOD):
        j = (j + S[i] + key[i% key_len]) % MOD
        S[i],S[j] = S[j],S[i]
    
    return S

'''
Psudo Random Generation Algorithm
'''
    
def PRGA(S):
    i = 0
    j = 0
    while True:
        i = (i + 1) % MOD
        j = (j + S[i]) % MOD
        
        S[i], S[j] = S[j], S[i]
        K = S[(S[i] + S[j]) % MOD]
        
        yield K

def get_keystream(key):
    
    S = KSA(key)
    return PRGA(S)


def encrypt_logic( key, text):
    
    key = [ ord(c) for c in key ]
    
    keystream = get_keystream(key)
    
    res = []
    for c in text:
        val = ("%02X" %(c ^ next(keystream)))
        res.append(val)
        
    return ''.join(res)

def encrypt( key, plaintext):
    plaintext = [ ord(c) for c in plaintext ]
    return encrypt_logic( key, plaintext)

def decrypt( key, ciphertext):
    ciphertext = codecs.decode( ciphertext, 'hex_codec')
    res = encrypt_logic( key, ciphertext)
    return codecs.decode(res, 'hex_codec').decode('utf-8')



def main():
    key = 'hello-fucking-python'
    plaintext = 'Life is a fucking movie,fuck you, fuck me, fuck everyone!'
    ciphertext = encrypt( key, plaintext)
    origtext = decrypt(key, ciphertext)
    print(plaintext == origtext)
    print(origtext)
    
def test():
    
    # Test case 1
    # key = 'Fuck'
    # plaintext = 'Fuck America'
    # ciphertext = '137021f3dc8343eef038241e'
    key = 'Fuck'
    ptext = 'Fuck America'
    ctext = encrypt( key, ptext)
    
    print(ctext)
    if ctext != "137021F3DC8343EEF038241E":
        print("fucking error: test case 1 encrypt error!!!")
    
    otext = decrypt( key, ctext)
    if otext != ptext:
        print("fucking error: test case 1 decrypt error!!!")
        
    # Text case 2
    # key = 'China'
    # plaintext = 'I love china!!!I love huawei!!!'
    # ciphertext= 'E655CD265A7E3288327367ED9B22850867058CB6FB64BB4872429731861815'
    key = 'China'
    ptext = 'I love china!!!I love huawei!!!'
    ctext = encrypt( key, ptext)
    
    print(ctext)
    if ctext != "E655CD265A7E3288327367ED9B22850867058CB6FB64BB4872429731861815":
        print("fucking error: test case 2 encrypt error!!!")
    
    otext = decrypt( key, ctext)
    if otext != ptext:
        print("fucking error: test case 2 decrypt error!!!")
    
    
    # Test case 3
    # key = 'Fruits'
    # plaintext = 'Orange and Banana and Watermelon'
    # ciphertext= 'A16B132F1F6CC8E20E4A7345D293A828773CA4E13F3D4217227F332002F03920'
    
    key = 'Fruits'
    ptext = 'Orange and Banana and Watermelon'
    ctext = encrypt( key, ptext)
    
    print(ctext)
    if ctext != 'A16B132F1F6CC8E20E4A7345D293A828773CA4E13F3D4217227F332002F03920':
        print("fucking error: test case 3 encrypt error!!!")
    
    otext = decrypt( key, ctext)
    if otext != ptext:
        print("fucking error: test case 3 decrypt error!!!")
    
    
main()